#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
HHL Phase-Estimation Interface + Pauli Trotter (Suzuki2) Evolution Oracle
=========================================================================

This module provides:
  1) A modular Hamiltonian simulator over Pauli terms with first-order and Suzuki2 steps.
  2) A factory that builds an EvolutionOracle(qc, ctrl, wrk, tau) from an arbitrary
     2^n x 2^n Hermitian matrix H. The oracle *includes the identity-term phase*
     correctly via a single-qubit RZ on the control (critical for HHL).
  3) Minimal, measurement-free PE encode/uncompute helpers that you can drop into HHL.

Usage (sketch):
    from hhl_pe_pauli_oracle import (
        build_pauli_evolution_oracle_from_matrix,
        hhl_pe_encode, hhl_pe_uncompute, pick_t_from_lambda_max
    )
    oracle = build_pauli_evolution_oracle_from_matrix(H, method='suzuki2', num_slices=4)
    t = pick_t_from_lambda_max(lambda_max_est, target_phase=0.25, safety=0.95)

    # Inside your HHL circuit:
    hhl_pe_encode(qc, phase_reg, work_reg, oracle, t)
    apply_reciprocal_rotation(qc, phase_reg, anc_flag, scale=2*math.pi/t)  # you implement
    hhl_pe_uncompute(qc, phase_reg, work_reg, oracle, t)
"""

from typing import Callable, List, Tuple, Dict, Optional
import math
import numpy as np
from numpy.linalg import eigvalsh

from qiskit import QuantumCircuit, QuantumRegister
from qiskit.circuit.library import QFT
from qiskit.circuit import Qubit

# -----------------------------
# Types
# -----------------------------

# evolution_oracle(qc, ctrl, wrk, tau): append controlled-U(tau) with control qubit `ctrl`
EvolutionOracle = Callable[[QuantumCircuit, Qubit, QuantumRegister, float], None]


# -----------------------------
# Pauli utilities and decomposition
# -----------------------------

PAULI_SINGLE = {
    'I': np.array([[1, 0],[0, 1]], dtype=complex),
    'X': np.array([[0, 1],[1, 0]], dtype=complex),
    'Y': np.array([[0, -1j],[1j, 0]], dtype=complex),
    'Z': np.array([[1, 0],[0, -1]], dtype=complex),
}

def all_pauli_strings(n: int) -> List[str]:
    out = ['I','X','Y','Z']
    for _ in range(n-1):
        out = [p+q for p in out for q in ['I','X','Y','Z']]
    return out

def pauli_matrix(pauli: str) -> np.ndarray:
    """
    LSB 映射：字符串第 0 位作用在 wrk[0] 上（右侧 Kronecker 因子）。
    例如 'ZX' -> np.kron(X, Z)；'XY' -> np.kron(Y, X)。
    """
    M = PAULI_SINGLE[pauli[0]]
    for ch in pauli[1:]:
        M = np.kron(PAULI_SINGLE[ch], M)
    return M

def decompose_to_paulis_with_identity(H: np.ndarray) -> Tuple[float, int, Dict[str, float]]:
    """Return (c_I, n_qubits, coeffs dict over non-identity Paulis)."""
    dim = H.shape[0]
    n = int(round(math.log2(dim)))
    if 2**n != dim:
        raise ValueError("H must be 2^n x 2^n.")
    c_I = float((np.trace(H).real) / (2**n))
    coeffs: Dict[str, float] = {}
    for p in all_pauli_strings(n):
        c = (np.trace(pauli_matrix(p).conj().T @ H) / (2**n)).real
        if p == 'I'*n:
            continue  # tracked as c_I
        if abs(c) > 1e-12:
            coeffs[p] = float(c)
    return c_I, n, coeffs


# -----------------------------
# Pauli-term Hamiltonian simulator (first-order / Suzuki2)
# -----------------------------

class HamiltonianSimulator:
    """Controlled exp(-i t H) over Pauli terms via Trotterization."""
    def __init__(self, terms: Dict[str, float], method: str = 'suzuki2', num_slices: int = 1):
        """
        terms: dict {pauli_string: coeff} excluding the identity term (handled separately)
        """
        if method not in ('first', 'suzuki2'):
            raise ValueError("method must be 'first' or 'suzuki2'")
        self.terms = terms
        self.method = method
        self.num_slices = max(1, int(num_slices))

    @staticmethod
    def _basis_change(qc: QuantumCircuit, qb, op: str):
        if op in ('I','Z'):
            return
        if op == 'X':
            qc.h(qb)
        elif op == 'Y':
            qc.sdg(qb); qc.h(qb)

    @staticmethod
    def _basis_uncompute(qc: QuantumCircuit, qb, op: str):
        if op in ('I','Z'):
            return
        if op == 'X':
            qc.h(qb)
        elif op == 'Y':
            qc.h(qb); qc.s(qb)

    @staticmethod
    def _add_ctrl_pauli_evolution(qc, ctrl, wrk, pauli, theta):
        """
        实现 e^{-i theta (⊗_i P_i)} 的受控版本（control=ctrl）。
        步骤：基变换到 Z -> CNOT 链归并到 anchor -> CRZ(2*theta) -> 撤销 CNOT -> 反基变换
        """
        ops = list(pauli)
        idx = [i for i, op in enumerate(ops) if op != 'I']
        if not idx:
            return  # 纯 I 项不在此处理（已在 factory 用 RZ(-c_I*tau) 加到 ctrl 上）

        # 1) 基变换：X->H，Y->S†;H
        for i in idx:
            HamiltonianSimulator._basis_change(qc, wrk[i], ops[i])

        # 2) 奇偶归并到锚点
        anchor = idx[-1]
        for j in idx[:-1]:
            qc.cx(wrk[j], wrk[anchor])

        # 3) 受控相位：exp(-i theta Z_anchor) 的受控实现 == CRZ(2*theta)
        qc.crz(2.0 * theta, ctrl, wrk[anchor])

        # 4) 撤销 CNOT 链
        for j in reversed(idx[:-1]):
            qc.cx(wrk[j], wrk[anchor])

        # 5) 反基变换
        for i in reversed(idx):
            HamiltonianSimulator._basis_uncompute(qc, wrk[i], ops[i])


    def _slice_first(self, qc: QuantumCircuit, ctrl: Qubit, wrk: QuantumRegister, dt: float):
        for p, c in self.terms.items():
            self._add_ctrl_pauli_evolution(qc, ctrl, wrk, p, theta=c * dt)

    def _slice_suzuki2(self, qc: QuantumCircuit, ctrl: Qubit, wrk: QuantumRegister, dt: float):
        half = 0.5 * dt
        items = list(self.terms.items())
        for p, c in items:
            self._add_ctrl_pauli_evolution(qc, ctrl, wrk, p, theta=c * half)
        for p, c in reversed(items):
            self._add_ctrl_pauli_evolution(qc, ctrl, wrk, p, theta=c * half)

    def add_controlled_time_evolution(self, qc: QuantumCircuit, ctrl: Qubit, wrk: QuantumRegister, total_time: float):
        dt = total_time / self.num_slices
        for _ in range(self.num_slices):
            if self.method == 'first':
                self._slice_first(qc, ctrl, wrk, dt)
            else:
                self._slice_suzuki2(qc, ctrl, wrk, dt)


# -----------------------------
# Evolution oracle factory (INCLUDES identity term phase)
# -----------------------------

def build_pauli_evolution_oracle_from_matrix(
    H: np.ndarray,
    method: str = 'suzuki2',
    num_slices: int = 2
) -> EvolutionOracle:
    """
    Build an EvolutionOracle for A=H that *includes* the identity term contribution.
    For a term c_I I, controlled-U(tau) adds a relative phase diag(1, e^{-i c_I tau})
    on the control, which equals (up to a global) an RZ(- c_I * tau) on the control qubit.
    """
    c_I, n, coeffs = decompose_to_paulis_with_identity(H)
    sim = HamiltonianSimulator(coeffs, method=method, num_slices=num_slices)

    def evolution_oracle(qc: QuantumCircuit, ctrl: Qubit, wrk: QuantumRegister, tau: float):
        # 1) Identity term as a single-qubit phase on control (global ignored):
        #    diag(1, e^{-i c_I tau}) = e^{-i c_I tau/2} * RZ(- c_I tau)
        if abs(c_I) > 1e-12:
            qc.rz(- c_I * tau, ctrl)

        # 2) Non-identity Pauli terms via Trotter
        sim.add_controlled_time_evolution(qc, ctrl, wrk, total_time=tau)

    return evolution_oracle

# test_hhl_pe_pauli_oracle.py
# -*- coding: utf-8 -*-

# import math
# import numpy as np
import pytest

# from qiskit import QuantumCircuit, QuantumRegister
from qiskit.quantum_info import Statevector, Operator

rng = np.random.default_rng(2025)


# -----------------------------
# 工具函数（不依赖 SciPy）
# -----------------------------

def expm_hermitian(H: np.ndarray, tau: float) -> np.ndarray:
    """对厄米矩阵的指数: exp(-i*tau*H)。用特征分解避免 SciPy 依赖。"""
    w, V = np.linalg.eigh(H)
    phase = np.exp(-1j * tau * w)
    return (V * phase) @ V.conj().T

def block_diag_control_U(U: np.ndarray) -> np.ndarray:
    """
    期望矩阵（控制位为 LSB，即右侧张量因子）:
        M = I_work ⊗ |0><0| + U_work ⊗ |1><1|
    这样就与 Qiskit 的 qubit-0=LSB 的矩阵顺序一致。
    """
    d = U.shape[0]
    P0 = np.array([[1, 0],
                   [0, 0]], dtype=complex)
    P1 = np.array([[0, 0],
                   [0, 1]], dtype=complex)
    return np.kron(np.eye(d, dtype=complex), P0) + np.kron(U, P1)


def circuit_unitary(qc: QuantumCircuit) -> np.ndarray:
    """拿到无测量电路的整体幺正矩阵。"""
    return Operator(qc).data

def unitary_global_distance(U: np.ndarray, V: np.ndarray) -> float:
    """
    全局相位不敏感的 Frobenius 距离：
    min_phi || U - e^{i phi} V ||_F / U.shape[0]
    """
    tr = np.trace(V.conj().T @ U)
    phi = np.angle(tr)  # 使 trace 最大的相位
    return np.linalg.norm(U - np.exp(1j * phi) * V, ord='fro') / U.shape[0]

def random_pauli_sum(n: int, k_terms: int = None, coef_scale: float = 1.0):
    """在 n 比特上随机挑若干个非 I 的 Pauli 串及实系数，返回 (c_I, dict)。"""
    if k_terms is None:
        k_terms = max(1, n + 1)
    candidates = [p for p in all_pauli_strings(n) if set(p) != {'I'}]
    rng.shuffle(candidates)
    chosen = candidates[:k_terms]
    coeffs = {p: float(coef_scale * rng.normal()) for p in chosen}
    c_I = float(coef_scale * rng.normal())
    return c_I, coeffs


# -----------------------------
# 1) Pauli 分解正确性（含 c_I）
# -----------------------------

@pytest.mark.parametrize("n", [1, 2])
def test_pauli_decomposition_roundtrip(n):
    # 构造 H = c_I I + sum c_p P
    c_I, coeffs = random_pauli_sum(n, k_terms=2*n+1, coef_scale=0.7)
    H = c_I * np.eye(2**n, dtype=complex)
    for p, c in coeffs.items():
        H = H + c * pauli_matrix(p)

    c_I2, n2, coeffs2 = decompose_to_paulis_with_identity(H)
    assert n2 == n, f"decompose_to_paulis_with_identity 返回的 n={n2} 与期待的 {n} 不符。"

    H_rec = c_I2 * np.eye(2**n, dtype=complex)
    for p, c in coeffs2.items():
        H_rec = H_rec + c * pauli_matrix(p)

    err = np.linalg.norm(H - H_rec, ord='fro') / np.linalg.norm(H, ord='fro')
    assert err < 1e-10, f"Pauli 分解/重构误差过大: {err:.3e}"


# -----------------------------
# 2) 基变换回滚应为恒等
# -----------------------------

@pytest.mark.parametrize("op", list("IXYZ"))
def test_basis_change_uncompute_is_identity(op):
    wrk = QuantumRegister(1, "w")
    qc = QuantumCircuit(wrk)
    HamiltonianSimulator._basis_change(qc, wrk[0], op)
    HamiltonianSimulator._basis_uncompute(qc, wrk[0], op)
    U = circuit_unitary(qc)
    err = unitary_global_distance(U, np.eye(2))
    assert err < 1e-12, f"基变换+反变换未还原恒等 (op={op})，误差={err:.3e}"


# -----------------------------
# 3) 受控单 Pauli 演化应等于精确受控 U
# -----------------------------

@pytest.mark.parametrize("pauli", ["X", "Y", "Z", "XX", "XY", "XZ", "YZ", "ZX"])
def test_ctrl_single_pauli_matches_exact(pauli):
    n = len(pauli)
    c = 0.7
    tau = 0.53

    # 目标 U = e^{-i tau * c * P}
    P = pauli_matrix(pauli)
    U = expm_hermitian(c * P, tau)

    # 电路：控制位 + n 工作位，应用一片 Pauli 受控演化（应当“精确”）
    ctrl = QuantumRegister(1, "c")
    wrk = QuantumRegister(n, "w")
    qc = QuantumCircuit(ctrl, wrk)
    HamiltonianSimulator({pauli: c}, method='suzuki2', num_slices=1)._add_ctrl_pauli_evolution(
        qc, ctrl[0], wrk, pauli, theta=c * tau
    )
    U_circ = circuit_unitary(qc)

    U_target = block_diag_control_U(U)
    err = unitary_global_distance(U_circ, U_target)
    assert err < 1e-10, f"受控单 Pauli 项不匹配 (P={pauli})，误差={err:.3e}"


# -----------------------------
# 4) 仅有恒等项：应等价于控制位 RZ(-c_I*tau)
# -----------------------------

@pytest.mark.parametrize("n", [1, 2, 3])
def test_identity_only_oracle(n):
    c_I = 1.234
    tau = 0.789

    H = c_I * np.eye(2**n, dtype=complex)
    oracle = build_pauli_evolution_oracle_from_matrix(H, method='suzuki2', num_slices=1)

    ctrl = QuantumRegister(1, "c")
    wrk = QuantumRegister(n, "w")
    qc = QuantumCircuit(ctrl, wrk)
    oracle(qc, ctrl[0], wrk, tau)
    U_circ = circuit_unitary(qc)

    # 目标：diag(I, e^{-i c_I tau} I)
    U = np.exp(-1j * c_I * tau) * np.eye(2**n, dtype=complex)
    U_target = block_diag_control_U(U)

    err = unitary_global_distance(U_circ, U_target)
    assert err < 1e-12, f"恒等项相位实现不正确，误差={err:.3e}。检查 qc.rz(-c_I*tau, ctrl) 符号或位置。"


# -----------------------------
# 5) 工厂：单一非 I Pauli 项应精确
# -----------------------------

@pytest.mark.parametrize("pauli", ["Z", "X", "Y", "ZX", "XZ"])
def test_factory_single_term_exact(pauli):
    n = len(pauli)
    c = -0.9
    tau = 0.37
    H = c * pauli_matrix(pauli)

    oracle = build_pauli_evolution_oracle_from_matrix(H, method='suzuki2', num_slices=1)

    ctrl = QuantumRegister(1, "c")
    wrk = QuantumRegister(n, "w")
    qc = QuantumCircuit(ctrl, wrk)
    oracle(qc, ctrl[0], wrk, tau)
    U_circ = circuit_unitary(qc)

    U = expm_hermitian(H, tau)
    U_target = block_diag_control_U(U)
    err = unitary_global_distance(U_circ, U_target)
    assert err < 1e-10, f"工厂生成的受控单项演化不精确 (P={pauli})，误差={err:.3e}"


# -----------------------------
# 6) 非对易项：Suzuki2 应优于一阶
# -----------------------------

def test_suzuki2_better_than_first_order_on_noncommuting():
    # 构造一个明显非对易的 H
    n = 2
    H = 0.8 * pauli_matrix("XZ") + 0.6 * pauli_matrix("ZX") + 0.3 * pauli_matrix("YI")
    tau = 0.8

    # 目标
    U = expm_hermitian(H, tau)
    U_target = block_diag_control_U(U)

    # 一阶
    oracle_first = build_pauli_evolution_oracle_from_matrix(H, method='first', num_slices=1)
    ctrl = QuantumRegister(1, "c")
    wrk = QuantumRegister(n, "w")
    qc1 = QuantumCircuit(ctrl, wrk)
    oracle_first(qc1, ctrl[0], wrk, tau)
    U1 = circuit_unitary(qc1)
    err_first = unitary_global_distance(U1, U_target)

    # Suzuki2
    oracle_s2 = build_pauli_evolution_oracle_from_matrix(H, method='suzuki2', num_slices=1)
    ctrl = QuantumRegister(1, "c")
    wrk = QuantumRegister(n, "w")
    qc2 = QuantumCircuit(ctrl, wrk)
    oracle_s2(qc2, ctrl[0], wrk, tau)
    U2 = circuit_unitary(qc2)
    err_s2 = unitary_global_distance(U2, U_target)

    assert err_s2 + 1e-12 <= err_first, (
        f"Suzuki2 并未优于一阶：err_first={err_first:.3e}, err_s2={err_s2:.3e}。"
        "若偶发失败，可增大 tau 或改用更强非对易的 H。"
    )


# -----------------------------
# 7) Suzuki2 随切片数增大，误差应下降
# -----------------------------

def test_suzuki2_converges_with_more_slices():
    n = 2
    # 人为设置非对易项
    H = 1.1 * pauli_matrix("XY") + 0.9 * pauli_matrix("YZ") + 0.7 * pauli_matrix("ZX")
    tau = 1.2
    U_target = block_diag_control_U(expm_hermitian(H, tau))

    errs = []
    for slices in [1, 2, 4, 8]:
        oracle = build_pauli_evolution_oracle_from_matrix(H, method='suzuki2', num_slices=slices)
        ctrl = QuantumRegister(1, "c")
        wrk = QuantumRegister(n, "w")
        qc = QuantumCircuit(ctrl, wrk)
        oracle(qc, ctrl[0], wrk, tau)
        U_circ = circuit_unitary(qc)
        err = unitary_global_distance(U_circ, U_target)
        errs.append(err)

    # 单调非增（允许浮点极小波动）
    for i in range(1, len(errs)):
        assert errs[i] <= errs[i-1] + 5e-13, f"Suzuki2 未随切片数增大而收敛: {errs}"


# -----------------------------
# 8) 工厂 + 随机厄米：小步长应高精度
# -----------------------------

@pytest.mark.parametrize("n", [1, 2])
def test_random_hermitian_small_tau_matches_well(n):
    # 用随机 Pauli 线性组合构造 H（含 c_I）
    c_I, coeffs = random_pauli_sum(n, k_terms=2*n+1, coef_scale=0.5)
    H = c_I * np.eye(2**n, dtype=complex)
    for p, c in coeffs.items():
        H = H + c * pauli_matrix(p)

    tau = 0.03  # 小步长
    oracle = build_pauli_evolution_oracle_from_matrix(H, method='suzuki2', num_slices=2)

    ctrl = QuantumRegister(1, "c")
    wrk = QuantumRegister(n, "w")
    qc = QuantumCircuit(ctrl, wrk)
    oracle(qc, ctrl[0], wrk, tau)
    U_circ = circuit_unitary(qc)

    U = expm_hermitian(H, tau)
    U_target = block_diag_control_U(U)
    err = unitary_global_distance(U_circ, U_target)

    assert err < 5e-5, f"随机厄米 + 小步长误差偏大: {err:.3e}。可能基变换/CNOT 链/CRZ 角度有问题。"


# -----------------------------
# 9) 签名/接口形状检查
# -----------------------------

def test_oracle_signature_shape():
    n = 2
    H = rng.normal(size=(2**n, 2**n)) + 1j * rng.normal(size=(2**n, 2**n))
    H = 0.5 * (H + H.conj().T)  # 厄米化
    oracle = build_pauli_evolution_oracle_from_matrix(H, method='suzuki2', num_slices=3)

    ctrl = QuantumRegister(1, "c")
    wrk = QuantumRegister(n, "w")
    qc = QuantumCircuit(ctrl, wrk)

    # 能否正常调用
    try:
        oracle(qc, ctrl[0], wrk, 0.123)
    except Exception as e:
        pytest.fail(f"EvolutionOracle 调用失败：{e!r}")